import torch
from MyData import MyDataset
from torch.utils.data import DataLoader
from net import Model
# Adam 和 AdamW 是两种常用的优化器，它们都是基于梯度下降算法的优化器，用于训练神经网络。它们之间的主要区别在于权重衰减（Weight Decay）的处理方式
from transformers import BertTokenizer
# Transformers 4.50.3 已移除 optimization.AdamW 的显式导出，或需依赖 PyTorch 的 torch.optim.AdamW 实现‌
from torch.optim import AdamW

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 30000

# 分词器
token = BertTokenizer.from_pretrained(
    r"../my_model_cache/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")


"""
将一批原始文本数据统一编码为BERT需要的格式，保证每个batch的数据长度一致
"""
def collate_fn(data):
    # 提取每个样本的文本数据
    sentes = [i[0] for i in data]
    # 提取每个样本的标签
    label = [i[1] for i in data]
    # print(sentes)
    # 使用BERT分词器对文本数据进行编码，返回一个包含input_ids、attention_mask、token_type_ids等的字典
    data = token.batch_encode_plus(batch_text_or_text_pairs=sentes,
                                   truncation=True,
                                   padding="max_length",
                                   max_length=500,
                                   return_tensors="pt",
                                   return_length=True)
    # print(data)
    input_ids = data["input_ids"]
    attention_mask = data["attention_mask"]
    token_type_ids = data["token_type_ids"]
    labels = torch.LongTensor(label)
    # print(input_ids,attention_mask,token_type_ids)
    return input_ids, attention_mask, token_type_ids, labels


# 创建数据集
train_dataset = MyDataset("train")
# 创建dataloader
train_laoder = DataLoader(dataset=train_dataset,
                          batch_size=100,  # 每批100个样本
                          shuffle=True,  # 打乱数据顺序
                          drop_last=True,  # 丢弃最后不足一个batch的数据
                          collate_fn=collate_fn  # 使用自定义的批处理函数
                          )

if __name__ == '__main__':
    # 开始训练
    print(DEVICE)
    model = Model().to(DEVICE)
    # 使用AdamW优化器优化模型参数，学习率为5e-4
    optimizer = AdamW(model.parameters(), lr=5e-4)
    # 使用交叉熵损失函数计算损失
    loss_func = torch.nn.CrossEntropyLoss()
    # 切换模型为训练模式（影响Dropout/BatchNorm等层）
    model.train()
    for epoch in range(EPOCH):
        # 遍历每个批次的数据
        for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_laoder):
            input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(
                DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
            # 使用模型前向传播计算输出
            out = model(input_ids, attention_mask, token_type_ids)
            # 计算模型输出和真实标签之间的损失
            loss = loss_func(out, labels)
            # 清零梯度
            optimizer.zero_grad()
            # 反向传播计算梯度
            loss.backward()
            # 更新模型参数
            optimizer.step()

            if i % 5 == 0:  # 每5个batch打印一次指标
                out = out.argmax(dim=1)   # 取概率最大的类别作为预测结果
                acc = (out == labels).sum().item() / len(labels)  # 计算准确率
                print(epoch, i, loss.item(), acc)
        torch.save(model.state_dict(), f"params/{epoch}bert01.pth")
        print(epoch, "参数保存成功！")
